from tqdm import tqdm
import os
import torch
from loader import loader4_baidu_senti
import run4_baidu_senti
import run6_train_interpret
# from nezha.modeling import nezha_senti
from nezha.modeling import nezha_interpret
from fish_tool import logs, sys_tool
from fish_tool.ai import torch_tool
from fish_tool.db.zip_tool import ZipTool
from loader.vocab import Vocab

# config = run4_baidu_senti.TrainConfig()
config = run6_train_interpret.TrainConfig()


def predict(interpret_tag, save_dir, show=False):
    vocab = Vocab(config.pre_model_dir)
    txt_path = os.path.join(save_dir, 'senti-rationale.txt')
    txt_f = open(txt_path, 'w', encoding='utf8')
    data = loader4_baidu_senti.get_data(config, 'test', batch_size=1)
    # model = nezha_senti.NeZhaForSequenceClassification(config)
    model = nezha_interpret.NeZhaForSequenceClassification(config)
    model = torch_tool.load_cpu_model(model, os.path.join(config.out_dir, 'model.pth'))
    model = torch_tool.cuda(model)
    model.eval()
    if not show:
        data = tqdm(data, desc='predict')
    is_test_start = False
    with torch.no_grad():
        for batch in data:
            # if 2178 in batch['data_ids']:
            #     is_test_start = True
            # if not is_test_start:
            #     continue
            resp = model(**batch, interpret_tag=interpret_tag, show=show)
            for info in resp['interprets']:
                txt = ','.join([str(t) for t in info['rationale']])
                info.setdefault('id', '--')
                line = '{id}\t{label}\t{txt}\n'.format(**info, txt=txt)
                txt_f.write(line)
    txt_f.close()
    txt = open(txt_path, encoding='utf8').read()
    zip_path = os.path.join(save_dir, 'senti-rationale.txt.zip')
    zip = ZipTool(zip_path)
    zip.add('senti-rationale.txt', txt)


if __name__ == '__main__':
    tag = 3
    save_dir = f'E:/code/learn/AI比赛/2021_12_山东_文本分类/2021-shandong-Smart-Grid-Classification/data/baidu_senti_20220514/{tag}'
    save_dir = f'/mnt/sda2/wangxiaoyu/tmp/nezha_pretrain_and_train/data/baidu_senti_20220514'
    os.makedirs(save_dir, exist_ok=True)
    # predict(tag, save_dir, show=True)
    predict(tag, save_dir, show=True)
